/***************************************************************************\
 * The Mattise Kernel														*
 * Copyright 2007 Matthew Iselin											*
 * Licensed under the GPL													*
 *																			*
 * mt.cc																	*
 *																			*
 * Handles task creation, and also handles virtual mode interrupt calls.	*
 * Utility functions for common tasks.										*
 *																			*
\***************************************************************************/

// necessary includes
#include <system.h>
#include <memory.h>
#include <console.h>
#include <disk.h>
#include <task.h>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <paging.h>
#include <pm.h>

// task states
#define TASK_MSGWAIT	3
#define TASK_ASLEEP		2
#define TASK_ALIVE		1 // ie. if( process.status ) is true if alive
#define TASK_DEAD		0

/// Performs a context switch
extern "C" void i386_context_switch(void **from_esp, void *esp);

/// Process table
//CLinkedList<struct tagStruct> ptable;
struct tagStruct ptable[1024]; // TODO: convert to linked list

/// Current task pointer
struct tagStruct current_task[1];

/// Current pid
int pid = -1;

/// Gets the current pid
int currpid()
{
	return pid;
}

/// Is this the first reschedule?
int firstr = 1;

/// Multitasking enabled?
bool enabled = false;

/// Gets the enabled state
bool mtenabled() { return enabled; };

/// Sets the enabled state to false
void mtdisable() { enabled = false; };

/// Sets the enabled state to true
void mtenable() { enabled = true; };

/// Initializes the list
void initmt()
{
	// fill in the process table
	for( int i = 0; i < 1024; i++ )
		ptable[i].taskState = TASK_DEAD;
	
	// we'll be switching this around before tasks are created
	// so make sure CR3 has something worth switching for
	current_task->cr3 = GetKernelPageDir();
}

/// Kills a process
void kill( int id )
{
	// take the process out of the list
	ptable[id].taskState = TASK_DEAD;

	// reschedule, so that we don't return to a dead task
	resched();
}

/// Waits for a process to finish execution
void WaitUntilComplete( int id )
{
	// wait until it's TASK_DEAD
	while( ptable[id].taskState != TASK_DEAD );
}

/// Sets a process' registers
void SetRegs( int p, struct regs* r )
{
	ptable[p].r.eax = r->eax;
	ptable[p].r.ebx = r->ebx;
	ptable[p].r.ecx = r->ecx;
	ptable[p].r.edx = r->edx;
	ptable[p].r.ebp = r->ebp;
	ptable[p].r.esp = r->esp;
	ptable[p].r.edi = r->edi;
	ptable[p].r.esi = r->esi;
}

/// Gets a process' registers
void GetRegs( int p, struct regs* r )
{
	r->eax = ptable[p].r.eax;
	r->ebx = ptable[p].r.ebx;
	r->ecx = ptable[p].r.ecx;
	r->edx = ptable[p].r.edx;
	r->ebp = ptable[p].r.ebp;
	r->esp = ptable[p].r.esp;
	r->esi = ptable[p].r.esi;
	r->edi = ptable[p].r.edi;
}

/// Executes a real-mode interrupt
void int_v86( uchar_t int_num, struct regs* r )
{
	// create space for it
	uchar_t* exc = (uchar_t*) 0x1000; /** TODO: dynamic allocation of space between 0x500 and 0x7C00 **/
	
	// fill in the necessary info
	*exc++ = 0xCD;
	*exc++ = int_num;
	*exc++ = 0xCD;
	*exc++ = 0x80; /** INT80 = kill task **/
	*exc++ = 0xEB;
	*exc++ = 0xFE;
	
	// create the task and wait for it to complete
	int p = CreateThread16bit( 0x1000, 1024, "Virtual 8086 Mode Task", r );
	
	// wait for it
	WaitUntilComplete( p );
	
	// done, so get the registers at time of ending
	r->eax = ptable[p].r.eax;
	r->ebx = ptable[p].r.ebx;
	r->ecx = ptable[p].r.ecx;
	r->edx = ptable[p].r.edx;
	r->esp = ptable[p].r.esp;
	r->ebp = ptable[p].r.ebp;
	r->esi = ptable[p].r.esi;
	r->edi = ptable[p].r.edi;
}

/// Called when a process exits
void die()
{
	// take the process out of the list
	ptable[pid].taskState = TASK_DEAD;
	
	// eventually we will just not execute this
	while( true );
}

/// This is a temporary storage space for the first context switch
struct tagStruct tmptask;

/// Dumps a stack
extern "C" void dump_stack( uint_t* esp )
{
	// dump the general purpose registers (meant to be at the bottom)
	kprintf( "Dumping stack...\nGeneral Purpose Registers:\n" );
	int i;
	for( i = 0; i < 8; i++ )
		kprintf( "esp[%d]: 0x%x\n", i, esp[i] );

	// dump the EIP
	kprintf( "EIP: 0x%x\n", esp[i++] );

	// dump CS
	//kprintf( "CS: 0x%x\n", esp[i++] );

	// dump EFLAGS
	//kprintf( "EFLAGS: 0x%x\n", esp[i++] );
}

/// Scheduler
void resched()
{
	// enabled?
	if( !enabled )
	{
		kprintf( "[MTSK] Multitasking disabled but resched() called.\n" );
		return;
	}
	
	// check that we can save - if so, save
	if( pid != -1 )
	{
		// do nothing, we'll fall through
	}
	else
	{
		if( ptable[0].taskState == TASK_ALIVE )
		{
			// perform the switch
			pid = 0; // new pid = 0
			asm volatile ( "movl %0,%%esp\r\npopa\r\nret" : : "r" ((uint_t)ptable[pid].esp) );
			kprintf( "FATAL ERROR: initial stack switch didn't work!\n" );
			while( 1 );
		}
		return;
	}
	
	// reselection cutoff
	ReselCutoff:

	kprintf( "resched\n" );
	
	// check that the process' timeslice is done
	if( ptable[pid].time <= 0 )
	{
		kprintf( "switching...\n" );

		// find the next alive task
		int i;
		for( i = pid+1; i < 1024; i++ )
		{
			if( ptable[i].taskState == TASK_ALIVE )
			{
				break;
			}
		}
		
		// check that i is valid
		if( i == 1024 )
		{
			i = 0;
			
			if( ptable[i].taskState == TASK_DEAD )
			{
				return;
			}
		}
		
		// set the time
		ptable[i].time = ptable[i].prio;

		kprintf( "Switching from %d to %d\n", pid, i );
		
		// set the new pid
		int oldpid = pid;
		int newpid = i;
		pid = newpid;

		// dump the new stack
		dump_stack( (uint_t*) ptable[i].esp );

		// perform the switch
		i386_context_switch( &ptable[oldpid].esp, ptable[newpid].esp );
	}
	else
	{
		// see if it got killed
		if( ptable[pid].taskState == TASK_DEAD )
		{
			// clear the time
			ptable[pid].time = 0;
			
			// i really dislike using goto but i cannot think of a neat way to do this)
			goto ReselCutoff;
		}
		
		// decrease the timeslice
		ptable[pid].time = 0; //ptable[pid].time--;
		goto ReselCutoff;
		
		// increase the cumulative time
		ptable[pid].ctime++;
	}
	
	// set the TSS entries
	global_tss.esp0 = (uint_t) (kstack + 1024);
	global_tss.ss0 = 0x10;
}

/// Task creation intro
int thread_prolog( uint_t stack, const char* pname, uint_t prio )
{	
	// find a spare entry
	int i;
	for( i = 0; i < 1024; i++ )
	{
		if( ptable[i].taskState == TASK_DEAD )
			break;
	}

	// check for overflow
	if( i == 1024 )
	{
		// none found
		kprintf( "\n[MTSK] Couldn't create process, too many processes!\n" );
		
		// die
		return -1;
	}
	
	// fill the entry
	ptable[i].taskState = TASK_ALIVE;
	ptable[i].prio = prio;
	ptable[i].time = ptable[i].ctime = 0;
	ptable[i].cr3 = GetKernelPageDir();
	memcpy( ptable[i].pname, pname, strlen( pname ) );
	void* dat = (void*) pm_alloc( stack );
	ptable[i].esp = (uint_t*) ((uint_t) dat) + (stack*0x1000);
	
	// return the index
	return i;
}

/// Creates a 16 bit task
int CreateThread16bit( uint_t base, uint_t stack, const char* pname, struct regs* r, uint_t prio )
{
	// get a process
	int i = thread_prolog( stack, pname, prio );
	
	// check
	if( i == -1 )
	{
		return -1;
	}

	// fill the stack
	uint_t* esp = (uint_t*) 0x2800;
	uint_t* ustack = (uint_t*) 0x2900;
	
	// entry and the id
	esp--; *esp = i; esp--;
	
	// return point
	*esp-- = (uint_t) die;
	
	// v86 segments
	*esp-- = 0x20;
	*esp-- = 0x20;
	*esp-- = 0x20;
	*esp-- = 0x20;
	
	// base of the task - 16-bit
	*esp-- = 0x20;
	*esp-- = (uint_t) ustack;
	*esp-- = 0x20202;
	*esp-- = 0;
	*esp-- = 0x1000;

	// error code and interrupt number
	*esp-- = 0; *esp-- = 0;
	
	// registers - 8 of them
	*esp-- = ptable[i].r.eax = r->eax;
	*esp-- = ptable[i].r.ecx = r->ecx;
	*esp-- = ptable[i].r.edx = r->edx;
	*esp-- = ptable[i].r.ebx = r->ebx;
	*esp-- = ptable[i].r.esp = r->esp;
	*esp-- = ptable[i].r.ebp = r->ebp;
	*esp-- = ptable[i].r.esi = r->esi;
	*esp-- = ptable[i].r.edi = r->edi;
	
	// KERNEL segments - 4 of them
/*	*esp-- = 0x10;
	*esp-- = 0x10;
	*esp-- = 0x10;
	*esp = 0x10;*/
	
	// save the pointer
	ptable[i].esp = esp;
	ptable[i].ustack = (uint_t) ustack;
	ptable[i].v8086 = true;
	ptable[i].base = 0x1000;
	
	// return the process's id
	return i;
}

/// Creates a new task (not 16-bit)
int CreateThread( uint_t base, uint_t stack, const char* pname, uint_t prio )
{
	// get a process entry
	int i = thread_prolog( stack, const_cast< char* >( pname ), prio );
	
	// check
	if( i == -1 )
	{
		return -1;
	}
	
	// save this pointer
	int topofstack = (int) ptable[i].esp;

	// fill the stack
	uint_t* esp = (uint_t*) ptable[i].esp;
	
	// entry and the id
	//esp--; *esp = i; esp--;
	
	// return point
//	*esp-- = (uint_t) die;
	
	// base of the task
/*	*esp-- = 0x202;
	*esp-- = 0x08;*/
	esp--;
	*esp = base;
	
	// error code and interrupt number
//	*esp-- = 0; *esp-- = 0;
	
	// registers - 8 of them
	for( int j = 0; j < 8; j++ )
	{
		esp--;
		*esp = 0;
	}
	
	// segments - 4 of them
/*	*esp-- = 0x10;
	*esp-- = 0x10;
	*esp-- = 0x10;
	*esp = 0x10;*/
	
	// save the pointer
	ptable[i].esp = esp;
	ptable[i].v8086 = false;
	ptable[i].base = base;
	
	// return the process's id
	return i;
}

/// Creates a new child thread
int CreateThread_UserSpace( int parent_pid, uint_t entrypoint, uint_t stack, const char* tname, uint_t prio )
{
	// get some space in the list
	int spc = thread_prolog( stack, const_cast< char* >( tname ), prio );
	
	// valid?
	if( spc == -1 )
		return -1;
	
	// the page directory for this thread
	uint_t pagedir = ptable[parent_pid].cr3;

	// calculate the virtual for the top of our stack
	uint_t virt_parent_ustack_top = ptable[parent_pid].ustack - (ptable[parent_pid].ustack_sz*0x1000);
	
	// our user stack bottom
	int thread_stack_bot = virt_parent_ustack_top - (stack*0x1000);
	
	// new stack goes into ustack
	ptable[parent_pid].ustack = thread_stack_bot - 0x1000;
	
	// allocate enough pages for the userstack and store them just below the page tables/directories
	uint_t userstack = (uint_t) pm_alloc( stack );
	uint_t isr_irq_return_stack = ((uint_t)ptable[spc].esp) - (stack*0x1000);
	for( int i = 0; i <= stack; i++ )
	{
		// ISR/IRQ return stack is mapped 1:1 to kernel and child process
		MapUserPhysicalAddress( isr_irq_return_stack + (i*0x1000), isr_irq_return_stack + (i*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );
		MapPhysicalAddress( pagedir, isr_irq_return_stack + (i*0x1000), isr_irq_return_stack + (i*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );

		// map the userstack
		MapPhysicalAddress( pagedir, thread_stack_bot + (i*0x1000), userstack + (i*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );
	}
	
	// setup the userstack
	uint_t taskring = 3;
	uint_t codeseg = 0x18;
	uint_t dataseg = 0x20;

	// fill the stack
	uint_t* esp = (uint_t*) ptable[spc].esp;
	
	// the thread id is at the top of the stack
	esp--; *esp = spc; esp--;

	// user stack
	*esp-- = dataseg | taskring;
	*esp-- = thread_stack_bot + (stack*0x1000);
	
	// base of the task - these guys run in ring 3
	*esp-- = 0x3202;
	*esp-- = codeseg | taskring;
	*esp-- = entrypoint;
	
	// error code and interrupt number
	*esp-- = 0; *esp-- = 0;
	
	// registers - 8 of them
	*esp-- = 0;
	*esp-- = 0; 
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	
	// segments - 4 of them
	*esp-- = dataseg | taskring;
	*esp-- = dataseg | taskring;
	*esp-- = dataseg | taskring;
	*esp = dataseg | taskring;
	
	// save the pointer
	ptable[spc].esp = esp;
	ptable[spc].v8086 = false;
	ptable[spc].base = entrypoint;
	ptable[spc].cr3 = pagedir;
	ptable[spc].taskState = TASK_ALIVE;
	ptable[spc].image_end = ptable[parent_pid].image_end;
	
	// return the threads's id
	return spc;
}

/// Stack top for all new processes
#define NEWPROCESS_STACKTOP		(0xD0000000)

/// Creates a new process, ie. one that sits in its own address space
int CreateProcess( uint_t pagedir, uint_t entrypoint, uint_t totalsize, uint_t stack, const char* pname, uint_t prio )
{
	// disable multitasking while we setup a new process
	mtdisable();
	
	// get a process entry - we can do this AS LONG AS
	// WE CHANGE THE CR3 VALUE LATER
	int i = thread_prolog( stack, const_cast< char* >( pname ), prio );
	
	// check
	if( i == -1 )
	{
		return -1;
	}
	
	// save this pointer
	int topofstack = (int) ptable[i].esp;
	
	// allocate space for the userstack
	uint_t userstack = (uint_t) pm_alloc( stack );

	// calculate the bottom of the user stack (virtual)
	uint_t ustack_bottom = NEWPROCESS_STACKTOP - (stack*0x1000);

	// physical location of the ISR/IRQ return stack
	uint_t isr_irq_return_stack = ((uint_t)ptable[i].esp) - (stack*0x1000);

	// pointer to the page directory
	uint_t* pdptr = (uint_t*) pagedir;

	// map in all the addresses we need
	for( int j = 0; j <= stack; j++ )
	{
		// ISR/IRQ return stack is mapped 1:1 to kernel and child process
		MapUserPhysicalAddress( isr_irq_return_stack + (j*0x1000), isr_irq_return_stack + (j*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );
		MapPhysicalAddress( pagedir, isr_irq_return_stack + (j*0x1000), isr_irq_return_stack + (j*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );

		// map the userstack
		MapPhysicalAddress( pagedir, ustack_bottom + (j*0x1000), userstack + (j*0x1000), PF_PRESENT | PF_READWRITE | PF_USER );
	}

	// save the userstack and its size into the process table
	ptable[i].ustack = NEWPROCESS_STACKTOP;
	ptable[i].ustack_sz = stack;
	
	// fill the stack
	uint_t* esp = (uint_t*) (((uint_t)ptable[i].esp) - 0x1000);
	
	// entry and the id
	esp--; *esp = i; esp--;
	
	// setup the userstack
	uint_t taskring = 3;
	uint_t codeseg = 0x18;
	uint_t dataseg = 0x20;
	
	// userstack
	*esp-- = dataseg | taskring;
	*esp-- = NEWPROCESS_STACKTOP;
	
	// base of the task - these guys run in ring 3, IOPL3
	*esp-- = 0x3202;
	*esp-- = codeseg | taskring;
	*esp-- = entrypoint;
	
	// error code and interrupt number
	*esp-- = 0; *esp-- = 0;
	
	// registers - 8 of them
	*esp-- = 0;
	*esp-- = 0; 
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	*esp-- = 0;
	
	// segments - 4 of them
	*esp-- = dataseg | taskring;
	*esp-- = dataseg | taskring;
	*esp-- = dataseg | taskring;
	*esp = dataseg | taskring;

	// save the pointer
	ptable[i].esp = esp;
	ptable[i].v8086 = false;
	ptable[i].base = entrypoint;
	ptable[i].cr3 = pagedir;
	ptable[i].taskState = TASK_ALIVE;
	uint_t imgend = totalsize + 4096;
	uint_t addend = imgend % 4096;
	ptable[i].image_end = imgend + addend;

	// enable multitasking again
	mtenable();
	
	// return the process's id
	return i;
}

// Note: the following two functions can be used to simulate a 'malloc'
// for userspace, without a free. Useful for making buffers within
// process memory. The sbrk() syscall uses these as well.

/// Gets the end of the image for a certain process (virtual address)
uint_t GetImageEnd( int id )
{
	return ptable[id].image_end;
}

/// Sets the end of the image for a certain process (virtual address)
uint_t SetImageEnd( int id, uint_t newend )
{
	return (ptable[id].image_end = newend);
}

/// Gets the CR3 for a process
int GetCR3( int id )
{
	return ptable[id].cr3;
}

/// Puts a message onto the list
void PutMessage( int id, struct tagMsg msg )
{
	// push the message onto that process's list
	ptable[id].msgList.Push( msg );
}

/// Gets a message from the list
void GetMessage( int id, struct tagMsg* msg )
{
	// take the first one
	*msg = ptable[id].msgList.Get( 0 );
	
	// delete it now
	ptable[id].msgList.Delete( 0 );
}

/// Gets the size of a message from the list
uint_t GetMessageSize( int id )
{
	// check for validity
	if( HasMessage( id ) == false )
		return -1;
	
	// take the first one
	struct tagMsg msg;
	msg = ptable[id].msgList.Get( 0 );
	
	// return the size
	return msg.msgLen;
}

/// Detects if a process has a message
bool HasMessage( int id )
{
	// size must not be 0
	return ( ( ptable[id].msgList.Size() != 0 ) ? true : false );
}

/// Gets the total time spent executing a task
long GetProcessTime( int id )
{
	// return it
	return ptable[id].ctime;
}

/// Resets the timeslice on a process (useful for yeilding control)
void ResetTimeslice( int id )
{
	// set it to 0
	ptable[id].time = 0;
}
